package com.hellohu.mbg.plugin;

import org.mybatis.generator.api.IntrospectedColumn;
import org.mybatis.generator.api.IntrospectedTable;
import org.mybatis.generator.api.PluginAdapter;
import org.mybatis.generator.api.dom.xml.Attribute;
import org.mybatis.generator.api.dom.xml.XmlElement;

import java.util.List;
import java.util.Properties;

/**
 * 在insert语句中增加
 * keyColumn="id" keyProperty="id" useGeneratedKeys="true"
 *
 * @author HuMingyang
 */
public class UseGeneratedKeysPlugin extends PluginAdapter {

    private static final String USE_GENERATED_KEYS = "useGeneratedKeys";
    private static final String KEY_COLUMN = "keyColumn";
    private static final String KEY_PROPERTY = "keyProperty";


    @Override
    public void setProperties(Properties properties) {
        super.setProperties(properties);
    }

    @Override
    public boolean validate(List<String> warnings) {
        return true;
    }

    @Override
    public boolean sqlMapInsertElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        genUseGeneratedKeys(element, introspectedTable);
        return true;
    }

    @Override
    public boolean sqlMapInsertSelectiveElementGenerated(XmlElement element, IntrospectedTable introspectedTable) {
        genUseGeneratedKeys(element, introspectedTable);
        return true;
    }

    private void genUseGeneratedKeys(XmlElement element, IntrospectedTable introspectedTable) {

        //通过introspectedTable的getPrimaryKeyColumns方法得到解析出来数据库中的主键列,且自增主键不会是复合主键
        List<IntrospectedColumn> keyColumns = introspectedTable
                .getPrimaryKeyColumns();
        if (keyColumns.size() != 1) {
            System.out.println(String.format("%s没有自增主键", introspectedTable.getTableConfiguration().getTableName()));
            return;
        }
        //得到这个唯一的主键列
        IntrospectedColumn keyColumn = keyColumns.get(0);
        //必须是递增主键
        if (!keyColumn.isAutoIncrement()) {
            System.out.println(String.format("%s没有自增主键", introspectedTable.getTableConfiguration().getTableName()));
            return;
        }
        //useGeneratedKeys
        element.addAttribute(new Attribute(USE_GENERATED_KEYS, "true"));
        //keyColumn
        element.addAttribute(new Attribute(KEY_COLUMN, keyColumn.getActualColumnName()));
        //keyProperty
        element.addAttribute(new Attribute(KEY_PROPERTY, keyColumn.getJavaProperty()));
    }
}
